open Facile
open Easy

type task = {
  id: Fd.t; (* id=n for NOP *)
  start_time: Fd.t;
  duration: Fd.t;
  end_time: Fd.t; 
}

let print_plan = fun plan ->
  Array.iter (fun t -> Fd.fprint stdout t.id; Printf.printf "\t") plan;
  print_newline ();
  Array.iter (fun t -> Fd.fprint stdout t.start_time; Printf.printf "\t") plan;
  print_newline();
  Array.iter (fun t -> Fd.fprint stdout t.duration; Printf.printf "\t") plan;
  print_newline()

(* constraint programming planner *)
let cplan = fun process_time transition deadline ->
  let n = Array.length process_time.(0)-1 and
      m = Array.length process_time in
  let plan = Array.init m (fun _ ->
    Array.init (n+1) (fun _ ->
      { id =    Fd.interval 0 n;
	start_time = Fd.interval 0 deadline;
	duration = Fd.interval 0 deadline;
	end_time = Fd.interval 0 deadline;
      })) in

  (* alloc.(i) = m means that i won't be performed by any agent *)
  let alloc = Array.init n (fun _ -> Fd.interval 0 m) in
  let nb_task = Array.init m 
    (fun k -> 
      Arith.e2fd
	(Arith.sum_fd
	    (Array.map
		(fun i -> Reify.boolean (fd2e i =~ i2e k)) alloc))) in


  (* flattened transition matrix *)
  let tr = Array.init m 
    (fun k -> Array.init ((n+1)*(n+1))
       (fun x ->
	  let t1 = x / (n+1) and t2 = x mod (n+1) in
	  Fd.interval transition.(k).(t1).(t2) transition.(k).(t1).(t2))) in

  let pt = Array.init m 
    (fun k -> Array.init (n+1)
      (fun i -> 
	Fd.interval process_time.(k).(i) process_time.(k).(i))) in

  let get_transition = fun k i j ->
    let t = Arith.e2fd ((fd2e i *~ i2e (n+1)) +~ fd2e j) in
    FdArray.get tr.(k) t in

  let makespan = 
    let c = Array.map (fun plan_k -> plan_k.(n).end_time) plan in
    FdArray.max c in

  let task_done = Arith.e2fd (Arith.sum_fd nb_task) in
  let task_undone = Arith.e2fd (i2e n -~ fd2e task_done) in

  let cost = Arith.e2fd ((i2e n -~ fd2e task_done) *~ i2e deadline +~ fd2e makespan) in

  (* cardinality check to ensure that each task is done only once *)
  let tasks = Array.init (m*(n+1)) 
    (fun x ->
      let k = x / (n+1) and t = x mod (n+1) in
      plan.(k).(t).id ) in

  let agent_tasks = Array.map (Array.map (fun t -> t.id)) plan in
  let agent_cc = Array.init m 
    (fun k -> Array.init (n+1) 
      (fun i -> 
	if i<n 
	then (Reify.boolean (fd2e alloc.(i) =~ i2e k), i)
	else ((Fd.interval 0 n), n)
      )) in

  Printf.printf "Init constraints\n%!";

  for k=0 to m-1 do
    Cstr.post (Gcc.cstr agent_tasks.(k) agent_cc.(k));

    (* init time *)
    Cstr.post (fd2e plan.(k).(0).start_time =~ i2e 0);
    Cstr.post (fd2e plan.(k).(0).duration =~ 
	(fd2e (FdArray.get pt.(k) plan.(k).(0).id) +~ 
	    fd2e 
	    (get_transition 
		k 
		(plan.(k).(n).id)
		(plan.(k).(0).id))
	));
  done;

  Printf.printf "State transition constraints\n%!";

  for k=0 to m-1 do
    for t=1 to n do
 
      (* ensure no idling between to steps *)
      Cstr.post (fd2e plan.(k).(t).start_time =~ fd2e plan.(k).(t-1).end_time);
 
      Cstr.post ((fd2e plan.(k).(t-1).id =~ i2e n) =>~~ 
	  (fd2e plan.(k).(t-1).id =~ (fd2e plan.(k).(t).id)));
 
      (* time constraint *)
      Cstr.post (fd2e plan.(k).(t).duration =~ 
	  fd2e (FdArray.get pt.(k) plan.(k).(t).id) +~ 
	  fd2e 
	  (get_transition 
	      k 
	      (plan.(k).(t-1).id)
	      (plan.(k).(t).id))
      );
 
      Cstr.post (fd2e plan.(k).(t).end_time =~ 
	  fd2e plan.(k).(t).start_time +~ fd2e plan.(k).(t).duration);
 
    done
  done;
 

  Printf.printf "Goals definition\n%!";

  let instantiate_hi v =
    Goals.instantiate (Domain.choose 
			  (fun x1 x2 -> x1 > x2)) v in
  let instantiate_med v = 
    Goals.instantiate (Domain.choose 
			  (fun x1 x2 -> abs (n/m - x1) < abs ( n/m - x2))) v in

  let instantiate_rand v =
    Goals.instantiate (Domain.choose
			  (fun x1 x2 -> Random.bool ())) v in

  let goal = 
    let ids = Array.map (Array.map (fun t -> t.id)) plan in
    let st = Array.map (Array.map (fun t -> t.start_time)) plan in
    let dr = Array.map (Array.map (fun t -> t.duration)) plan in
    
    let min_size_order = 
      Goals.Array.choose_index (fun a1 a2 -> Var.Attr.size a1 < Var.Attr.size a2)
    in

    (*(instantiate_hi task_done) &&~*)
    (Goals.Array.forall instantiate_med nb_task) &&~
      (Goals.Array.forall ~select: min_size_order Goals.indomain alloc) &&~
      (Goals.Array.forall 
	  (Goals.Array.forall ~select: min_size_order Goals.indomain) ids) &&~
      (Goals.Array.forall 
	  (Goals.Array.forall ~select: min_size_order Goals.indomain) st) &&~
      (Goals.Array.forall 
	  (Goals.Array.forall ~select: min_size_order Goals.indomain) dr)
  in

  let time = ref (Sys.time ()) in

  let print = Goals.atomic 
    (fun () -> 
      Printf.printf "%6.2f:\t%!" (Sys.time () -. !time);
      Fd.fprint stdout task_done; 
      Printf.printf "\t%!";
      Fd.fprint stdout makespan; 
      Printf.printf "\n%!"
    )
    in
  
  (*Goals.solve ((Goals.lds ~step:4 (goal &&~ print)) &&~ Goals.fail)*)

  Printf.printf "Search\n%!";
  
  Goals.solve 
    ((Goals.minimize 
	 ( (Goals.lds ~step: 4 goal) &&~ print) cost (fun _ -> ()) 
       &&~ Goals.fail) 
      ||~ Goals.success)


let main = fun () ->
  let n = (int_of_string Sys.argv.(1)) 
  and m = (int_of_string Sys.argv.(2)) in

  let width = 600 
  and height = 600 in
  let wind_dir = Random.float 360.
  and wind_speed = 5. *. (float (Random.int 4)) in
  let distance = fun (x1, y1) (x2, y2) -> 
    let x1 = 4. *. x1 and
	x2 = 4. *. x2 and
	y1 = 4. *. y1 and
	y2 = 4. *. y2 in
    sqrt ((x1-.x2)** 2. +. (y1-. y2)**2.) in
  let coord = 
    Array.init (n+1) 
      (fun k -> 
	if k < n
	then
	  ((float width) /. 2. -. Random.float (float width), 
	  (float height)/. 2. -. Random.float (float height))
	else
	  (0., 0.)
      ) in


  let task_time = 60 in
  let ac_speed = 20. in

  let pt = Array.init m 
    (fun k -> Array.init (n+1) (fun i -> if i<n then task_time else 0)) in
  let tr = Array.init m 
    (fun k -> Array.init (n+1) 
      (fun i -> Array.init (n+1)
	(fun j -> 
	  let alpha = 
	    let (x1, y1) = coord.(i) and
		(x2, y2) = coord.(j) in
	    (atan2 (y2 -. y1) (x2 -. x1)) -. wind_dir
	  in
	  let speed = wind_speed *. cos alpha 
	    +. sqrt (wind_speed**2. *. (sin alpha)**2. 
		   +. ac_speed**2.) in
	  int_of_float 
	    ((distance coord.(i) coord.(j)) /. speed)
	))) in

  let deadline = 10000 in

  (*
  for k=0 to m-1 do
    for i=0 to n do
      for j=0 to n do
      Printf.printf "%3d " tr.(k).(i).(j);
      done;
      Printf.printf "\n";
    done;
    Printf.printf "\n%!";
  done;
  print_newline();
  *)

  cplan pt tr deadline
    

let _ = main ()
